# Copyright (c) 2018 Presto Labs Pte. Ltd.
# Author: jaewon

import datetime
import logging
import pprint
import sys

import tabulate
from tornado import gen

from coin.base.datetime_util import to_datetime
from coin.base.timestamp import get_timestamp
from coin.exchange.base.order_gateway import OrderSide
from coin.strategy.mm.calculator import MovingSum
from coin.strategy.mm.bitmex.lazy_pull_mm_executor import (BitmexLazyPullMMExecutor)
from coin.strategy.mm.bitfinex.simple_mm_executor_2 import (BitfinexSimpleMMExecutor_v2)
from coin.strategy.mm.base.bitfinex_bitmex_strategy_base import (BitfinexCurrency,
                                                                 BitfinexProduct,
                                                                 BitmexFuturesProduct,
                                                                 StrategyBase)
from coin.strategy.mm.passive.base import (FillRow, FeatureRow)

# TODO(inkyu): Move me to coin.base?
from coin.exchange.base.tick import FixedTickPrice
from coin.exchange.base.tick import BitfinexBTCUSDTickPrice


def calculate_unrealized_pnl(pos, avg_price, ask0_price, bid0_price):
  if pos == 0:
    return 0
  mid_price = (ask0_price + bid0_price) / 2.
  return pos * (1. / avg_price - 1. / mid_price)


def get_side_sign(order_side):
  if order_side == OrderSide.BUY:
    return 1
  elif order_side == OrderSide.SELL:
    return -1
  raise ValueError('Invalid order side: %s' % order_side)


class PassivePassiveStrategy(StrategyBase):
  def __init__(self, strategy_params, dry_run=True):
    super().__init__(strategy_params['product_1'],
                     strategy_params['product_2'], ['BTC'],
                     logger=logging.getLogger('PassivePassiveStrategy'))
    self._strategy_params = strategy_params
    self._dry_run = dry_run

    self._product_1 = strategy_params['product_1']
    self._product_2 = strategy_params['product_2']
    self._tick_1 = strategy_params['tick_1']
    self._tick_2 = strategy_params['tick_2']

    # Pricing parameters
    self._basis_ma_window = (strategy_params['basis_ma_window_min'] * 60 * (10**9))
    self._true_book_1 = strategy_params['true_book_1']  # Bitfinex
    self._true_book_2 = strategy_params['true_book_2']  # Bitmex
    self._edge_1 = strategy_params['edge_bp_1'] / 10000.
    self._edge_2 = strategy_params['edge_bp_2'] / 10000.

    # Execution parameters
    self._lot_size_1 = strategy_params['lot_size_1']
    self._lot_size_2 = strategy_params['lot_size_2']
    self._reserve = strategy_params['reserve']
    self._imbalance_allowance = strategy_params['imbalance_allowance']

    # Bitfinex only
    self._base_pos_usd_1 = strategy_params['base_pos_usd_1']

    self._bs_pos = 0.
    self._sb_pos = 0.
    self._mid_basis_ma = MovingSum(self._basis_ma_window)

    self._executor_1 = None
    self._executor_2 = None

    self._last_feature_dump_ts = 0

    self._logger.info('Strategy params:\n%s', pprint.pformat(strategy_params))

    # Print header
    self._logger.info(FillRow.get_header())
    self._logger.info(FeatureRow.get_header())

  def on_base_init(self, ioloop, queue, flow_sub):
    super().on_base_init(ioloop, queue, flow_sub)

  def _get_pos_balance_1(self):
    # It's a negative position, in fact.
    usd_balance_1 = self._og_1.get_balance('USDT')
    pos_1 = self._base_pos_usd_1 - usd_balance_1
    balance_1 = self._og_1.get_balance('BTC')
    return pos_1, balance_1, usd_balance_1

  def _get_pos_balance_2(self):
    pos_2 = self._og_2.get_position(self._product_2)
    avg_entry_price_2 = self._og_2.get_average_entry_price(self._product_2)
    unrealized_pnl_2 = calculate_unrealized_pnl(pos_2,
                                                avg_entry_price_2,
                                                self._book_2.ask0().price,
                                                self._book_2.bid0().price)
    balance_2 = self._og_2.get_balance('BTC') + unrealized_pnl_2
    return pos_2, balance_2

  def _print_balance(self):  # Override the original version.
    if self._book_1 is None or self._book_2 is None:
      return

    pos_1, balance_1, usd_balance_1 = self._get_pos_balance_1()
    pos_2, balance_2 = self._get_pos_balance_2()

    midp_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
    corrected_balance_1 = balance_1 - pos_1 / midp_1

    table = [['*', 'BTC_Value', corrected_balance_1 + balance_2, corrected_balance_1,
              balance_2], ['*', 'BTC', balance_1 + balance_2, balance_1,
                           balance_2], ['*', 'USD', usd_balance_1, usd_balance_1,
                                        0], ['*', 'Position', pos_1 + pos_2, pos_1, pos_2],
             ['*', 'USD_BS', self._bs_pos, 0, 0], ['*', 'USD_SB', self._sb_pos, 0, 0]]

    self._logger.info(
        '----- Balance:\n* %s\n%s' %
        (str(datetime.datetime.now()),
         tabulate.tabulate(
             table, headers=['*', 'Currency', 'Total', 'Bitfinex', 'Bitmex'], floatfmt=".6f")))

  @gen.coroutine
  def init_executor(self):
    while self._book_1 is None or self._book_2 is None:
      yield gen.sleep(0.5)

    self._logger.info('Intializing passive executor...')

    # Bitfinex
    executor_params_1 = {
        'product': self._product_1,
        'lot_size': self._lot_size_1,  # USD
        'stickiness': 0.,
        'base_pos': self._base_pos_usd_1,
        'min_pos': -self._reserve,
        'max_pos': self._reserve,
        'order_update_period': 2.0,  # Seconds: Can be reducible
    }
    self._executor_1 = BitfinexSimpleMMExecutor_v2(executor_params_1, self._og_1)

    executor_params_2 = {
        'product': self._product_2,
        'lot_size': self._lot_size_2,
        'min_pos': -self._reserve,
        'max_pos': self._reserve,
        'order_update_period': 2.0,  # Second
        'post_only': True,
        'pull_threshold_bp': 0.
    }
    self._executor_2 = BitmexLazyPullMMExecutor(executor_params_2, self._og_2)

    # Register fill callback
    self._og_1.register_fill_callback(self.on_fill_1)
    self._og_2.register_fill_callback(self.on_fill_2)

    # Reset bs_pos, sb_pos
    pos_1, _balance_1, _usd_balance_1 = self._get_pos_balance_1()
    pos_2, _balance_2 = self._get_pos_balance_2()

    net_pos = pos_1 + pos_2
    self._bs_pos = net_pos / 2.
    self._sb_pos = net_pos / 2.

    self._logger.info('Initial BS=%.2f, SB=%.2f', self._bs_pos, self._sb_pos)

  def on_fill_1(self, fill_price, fill_qty_btc, order):
    # Bitfinex: qty [BTC]
    fill_qty_usd = fill_qty_btc * fill_price
    if order.order_side == OrderSide.BUY:
      self._bs_pos += fill_qty_usd
    elif order.order_side == OrderSide.SELL:
      self._sb_pos -= fill_qty_usd
    else:
      raise ValueError(order)
    self.dump_fill(1, fill_price, fill_qty_usd, order.order_side)

  def on_fill_2(self, fill_price, fill_qty_usd, order):
    # BitMEX: qty [USD]
    if order.order_side == OrderSide.BUY:
      self._sb_pos += fill_qty_usd
    elif order.order_side == OrderSide.SELL:
      self._bs_pos -= fill_qty_usd
    else:
      raise ValueError(order)
    self.dump_fill(2, fill_price, fill_qty_usd, order.order_side)

  def dump_fill(self, exchange_id, fill_price, fill_qty, order_side, timestamp=None):
    timestamp = timestamp or get_timestamp()

    pos_1, balance_1, _usd_balance_1 = self._get_pos_balance_1()
    if self._book_1 is not None:
      midp_1 = (self._book_1.ask0().price + self._book_1.bid0().price) / 2.
      corrected_balance_1 = balance_1 - pos_1 / midp_1
    else:
      corrected_balance_1 = balance_1 - pos_1 / 10000.

    pos_2, balance_2 = self._get_pos_balance_2()
    side_sign = get_side_sign(order_side)
    row = FillRow(timestamp,
                  to_datetime(timestamp),
                  exchange_id,
                  side_sign,
                  fill_price,
                  fill_qty,
                  pos_1,
                  pos_2,
                  self._bs_pos,
                  self._sb_pos,
                  corrected_balance_1,
                  balance_2)
    self._logger.info(row.to_log_string())

  def on_book(self):
    if self._book_1 is None or self._book_2 is None:
      return

    # Check book freshness
    book_ts = max(self._book_1.timestamp, self._book_2.timestamp)
    min_ts = min(self._book_1.timestamp, self._book_2.timestamp)
    if book_ts - min_ts > (2 * 10**9):
      return

    # Check book sanity
    if (self._book_1.ask0().price <= self._book_1.bid0().price
        or self._book_2.ask0().price <= self._book_2.bid0().price):
      return

    # Bitfinex
    ask0_1 = self._book_1.ask0().price
    bid0_1 = self._book_1.bid0().price
    askt_1 = self._book_1.get_notional_asks_by_amt(self._true_book_1)[1][-1][0]
    bidt_1 = self._book_1.get_notional_bids_by_amt(self._true_book_1)[1][-1][0]
    midt_1 = (askt_1 + bidt_1) / 2.

    # Bitmex
    ask0_2 = self._book_2.ask0().price
    bid0_2 = self._book_2.bid0().price
    askt_2 = self._book_2.get_notional_asks_by_qty(self._true_book_2)[1][-1][0]
    bidt_2 = self._book_2.get_notional_bids_by_qty(self._true_book_2)[1][-1][0]
    midt_2 = (askt_2 + bidt_2) / 2.

    mid_diff = midt_1 - midt_2
    self._mid_basis_ma.update(book_ts, mid_diff)

    if len(self._mid_basis_ma) < 2000:
      return

    tick_1 = self._tick_1
    tick_2 = self._tick_2
    edge_1 = self._edge_1
    edge_2 = self._edge_2

    midt_2_proj_1 = midt_2 + float(self._mid_basis_ma.average)
    midt_1_proj_2 = midt_1 - float(self._mid_basis_ma.average)

    ask0_minus_one_1 = tick_1.get_prev_price(ask0_1)
    bid0_plus_one_1 = tick_1.get_next_price(bid0_1)
    ask0_minus_one_2 = tick_2.get_prev_price(ask0_2)
    bid0_plus_one_2 = tick_2.get_next_price(bid0_2)

    askp_1 = max(tick_1.ceil(midt_2_proj_1 * (1. + edge_1 / 2.)), ask0_minus_one_1, bid0_plus_one_1)
    bidp_1 = min(tick_1.floor(midt_2_proj_1 * (1. - edge_1 / 2.)),
                 ask0_minus_one_1,
                 bid0_plus_one_1)
    askp_2 = max(tick_2.ceil(midt_1_proj_2 * (1. + edge_2 / 2.)), ask0_minus_one_2, bid0_plus_one_2)
    bidp_2 = min(tick_2.floor(midt_1_proj_2 * (1. - edge_2 / 2.)),
                 ask0_minus_one_2,
                 bid0_plus_one_2)

    if not self.is_og_ready() or self._dry_run:
      return

    self.post_1(askp_1, bidp_1)
    self.post_2(askp_2, bidp_2)

    now_ts = get_timestamp()
    if now_ts - self._last_feature_dump_ts >= 2 * (10**9):
      self._last_feature_dump_ts = now_ts
      row = FeatureRow(book_ts,
                       to_datetime(book_ts),
                       ask0_1,
                       bid0_1,
                       askt_1,
                       bidt_1,
                       midt_1,
                       ask0_2,
                       bid0_2,
                       askt_2,
                       bidt_2,
                       midt_2,
                       float(self._mid_basis_ma.average),
                       askp_1,
                       bidp_1,
                       askp_2,
                       bidp_2)
      self._logger.info(row.to_log_string())

  def post_1(self, askp, bidp):
    if self._sb_pos <= -self._imbalance_allowance:
      askp = None
    if self._bs_pos >= self._imbalance_allowance:
      bidp = None
    self._executor_1.manage_mm_orders(askp, bidp)

  def post_2(self, askp, bidp):
    if self._bs_pos <= -self._imbalance_allowance:
      askp = None
    if self._sb_pos >= self._imbalance_allowance:
      bidp = None
    self._executor_2.manage_mm_orders(askp, bidp)


def get_param_set(init_og_immediately=False):
  product_1 = BitfinexProduct.FromStr('BTC-USD')
  product_2 = BitmexFuturesProduct.FromStr('BTC-USD.PERPETUAL')

  edge_bp_1 = 20.  # Bitfinex
  edge_bp_2 = 20.  # Bitmex

  bitfinex_og_config_name = 'bitfinex_jshin0'
  bitmex_og_config_name = 'bitmex_xunke04'

  scale_factor = 5
  true_book_size = 10000.
  lot_size_usd = scale_factor * 1000.
  reserve_usd = lot_size_usd * 15.  # 75,000
  imbalance_allowance = lot_size_usd * 2.
  bitfinex_base_pos = reserve_usd + lot_size_usd * 3.5  # 92,500
  logging.info('bitfinex_base_pos=%.2f', bitfinex_base_pos)

  strategy_params = {
      'product_1': product_1,
      'product_2': product_2,
      'basis_ma_window_min': 30,
      'true_book_1': true_book_size,
      'true_book_2': true_book_size,
      'edge_bp_1': edge_bp_1,
      'edge_bp_2': edge_bp_2,
      'lot_size_1': lot_size_usd,
      'lot_size_2': lot_size_usd,
      'tick_1': BitfinexBTCUSDTickPrice(),
      'tick_2': FixedTickPrice(0.5),
      'reserve': reserve_usd,
      'base_pos_usd_1': bitfinex_base_pos,
      'imbalance_allowance': imbalance_allowance,
  }

  # Trading environment config
  if init_og_immediately:
    init_og_after_min = 2
  else:
    init_og_after_min = strategy_params['basis_ma_window_min']

  exit_after_min = init_og_after_min + 50
  env_config = {
      'first_currency': [
          BitfinexCurrency.FromStrCurrency('BTC'), BitfinexCurrency.FromStrCurrency('USDT')
      ],
      'first_product': product_1,
      'first_og_config_name': bitfinex_og_config_name,
      'second_product': product_2,
      'second_og_config_name': bitmex_og_config_name,
      'init_order_gateway_after': datetime.timedelta(minutes=init_og_after_min),
      'exit_after': datetime.timedelta(minutes=exit_after_min)
  }

  return strategy_params, env_config


def main(argv):
  strategy_params, env_config = get_param_set(init_og_immediately=False)

  # Run strategy!
  strategy = PassivePassiveStrategy(strategy_params, dry_run=False)
  return StrategyBase.run_live(on_base_init_callback=strategy.on_base_init,
                               on_book_reset_callback=strategy.on_book_reset,
                               on_order_gateway_reset_callback=strategy.on_og_reset,
                               **env_config)


if __name__ == '__main__':
  logging.basicConfig(level='DEBUG', format='%(levelname)8s %(asctime)s %(name)s] %(message)s')
  sys.exit(main(sys.argv))
